package translate

import (
	"context"
	"fmt"
	"github.com/openai/openai-go"
	"github.com/openai/openai-go/option"
	"log"
	"strings"
)

type TranslateResult struct {
	Code    string `json:"code"`
	Content string `json:"content"`
	Result  string `json:"result"`
}

func OpenAITranslate(url, apiKey, model, text, language, prompt string) (*TranslateResult, error) {
	log.Println("ai 翻译中...")
	client := openai.NewClient(
		option.WithBaseURL(url),
		option.WithAPIKey(apiKey),
	)
	res, err := client.Chat.Completions.New(context.TODO(), openai.ChatCompletionNewParams{
		Model: model,
		Messages: []openai.ChatCompletionMessageParamUnion{
			openai.SystemMessage(fmt.Sprintf(prompt, language)),
			openai.UserMessage(text),
		},
	})
	if err != nil {
		return nil, err
	}
	content := res.Choices[0].Message.Content

	log.Println("相应结果：", content)
	//处理相应think标签内容
	if strings.Index(content, "<think>") != -1 && strings.Index(content, "</think>") != -1 {
		content = strings.Split(content, "</think>")[1]
	}
	content = strings.Trim(content, "\n")

	return &TranslateResult{
		"0",
		text,
		content,
	}, nil
}
